"""
线性回归 on keras & sklearn
"""
import os
import numpy as np
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense
from keras import optimizers, losses, metrics
from keras import backend as K
from sklearn.metrics import r2_score
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import learning_curve
from sklearn.pipeline import Pipeline
from utils.plot_graph import keras_plot_model
from utils.model_info import ModelFitPredict


def load_data(w, b, size=10000, valid_ratio=0.2, test_ratio=0.2, dim=1):
    """
    生成数据
    :param w:
    :param b:
    :param size:
    :param valid_ratio:
    :param test_ratio:
    :param dim:
    :return:
    """

    x = np.linspace(-1, 1, num=size)
    np.random.shuffle(x)
    x = x.reshape((-1, dim))

    y = np.dot(x, w) + b
    noise = \
        np.random.normal(loc=0, scale=1, size=int(size / dim)).reshape(-1, 1)
    y += noise

    valid_size = int(size * valid_ratio)
    test_size = int(size * test_ratio)

    return x[: valid_size], y[: valid_size], \
           x[valid_size: valid_size + test_size], \
           y[valid_size: valid_size + test_size], \
           x[valid_size + test_size:], y[valid_size + test_size:], \
           x, y


def build_model(input_shape, output_dim=1, learning_rate=0.005):
    """

    :param input_shape: 特征数
    :param output_dim:
    :param learning_rate:
    :return:
    """
    model = Sequential()
    model.add(Dense(units=output_dim, input_dim=input_shape))
    # model.compile(optimizer=optimizers.SGD(learning_rate=learning_rate),
    #               loss=losses.mean_squared_error)
    model.compile(optimizer='sgd', loss='mse')
    return model


def sklearn_lr(x_train, y_train, x_test, y_test, x, y):
    """

    :param x_train:
    :param y_train:
    :param x_test:
    :param y_test:
    :param x:
    :param y:
    :return:
    """
    # 模型直接训练
    model = LinearRegression()

    # plot the graph of model's train & test score
    train_size, train_scores, test_scores = \
        learning_curve(estimator=model,
                       X=x, y=y,
                       train_sizes=np.linspace(0.1, 1.0, 10),
                       cv=10,
                       n_jobs=4)
    plt.plot(train_size, np.mean(train_scores, axis=1),
             'o-', c='r', label='train')
    plt.plot(train_size, np.mean(test_scores, axis=1),
             'o-', c='g', label='test')
    plt.legend(loc='best')
    plt.show()

    # fit
    model.fit(X=x_train, y=y_train)
    pred = model.predict(X=x_test)

    model_score = model.score(X=x_test, y=y_test)
    model_r2 = r2_score(y_true=y_test, y_pred=pred)
    weight = model.coef_
    b = model.intercept_

    # 管道
    pipe = Pipeline(steps=[('lr', LinearRegression())])
    pipe.fit(X=x_train, y=y_train)
    pip_pred = pipe.predict(X=x_test)

    pip_scores = pipe.score(x_test, y_test)
    pipe_r2 = r2_score(y_true=y_test, y_pred=pip_pred)

    model_info = 'weight: {}, b: {}, model score: {}, r2 score: {}'.format(
        weight, b, model_score, model_r2)

    pipe_info = 'model score: {}, r2 score: {}'.format(pip_scores, pipe_r2)

    print(model_info)
    print(pipe_info)


def lr_run(workers):
    """
    主程序
    :return:
    """
    w = np.array([[7.6], [-19.28]])
    # w = -19.7
    b = 23.3

    dim = len(w) if isinstance(w, np.ndarray) or isinstance(w, list) else 1

    x_train, y_train, x_valid, y_valid, x_test, y_test, x, y = \
        load_data(w=w, b=b, dim=dim)

    # 行数是样本量，列数是特征量(输入维度)
    input_dim = x_train.shape[1]
    output_dim = y_train.shape[1]

    # sklearn
    sklearn_lr(x_train=x_train, y_train=y_train,
               x_test=x_valid, y_test=y_valid,
               x=x, y=y)

    print('*' * 100)

    # kears
    model = build_model(input_shape=input_dim, output_dim=output_dim)

    # plot model
    keras_plot_model(model=model, pic='../pic/lr/线性回归模型.png')

    fit_predict = ModelFitPredict(model=model,
                                  train_x=x_train, train_y=y_train,
                                  valid_x=x_valid, valid_y=y_valid,
                                  test_x=x_test, test_y=y_test,
                                  epochs=20,
                                  model_name='线性回归',
                                  pic_path='../pic/lr',
                                  batch_size=32,
                                  workers=workers)

    pred_valid, pred_test = fit_predict.predict()
    valid_scores = r2_score(y_true=y_valid, y_pred=pred_valid)
    test_scores = r2_score(y_true=y_test, y_pred=pred_test)
    print('R^2得分：验证集： {}, 测试集： {}'.format(valid_scores, test_scores))
